package assignment1;

import java.io.File;
import java.io.IOException;
import java.io.PrintWriter;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import java.util.TreeMap;

/**
 * 
 * @author Kha Nguyen - 598240
 *
 */
public class ShakespeareWordCountApp {

    private final int numberOfThreads;

    private final String fullTextFile;

    private final String inputFile;

    private final String outputFileName;

    private final Thread[] workerThreads;

    private final CountWordTask[] workerTasks;

    public ShakespeareWordCountApp(int numberOfThreads, String fullTextFile, String inputFile, String outputFileName) {
        super();
        this.numberOfThreads = numberOfThreads;
        this.fullTextFile = fullTextFile;
        this.inputFile = inputFile;
        this.outputFileName = outputFileName;
        this.workerTasks = new CountWordTask[numberOfThreads];
        this.workerThreads = new Thread[numberOfThreads];
    }

    public void run() {
        final List<String> tokensToScan = getTextTokens(fullTextFile);
        final Set<String> uniqueTokens = getTokensToCount(inputFile);
        spawnAndRunWorkerThreads(tokensToScan, uniqueTokens);
        writeFinalResult();
    }

    private List<Set<String>> partitionSet(Set<String> setToPartition, int partitions) {
        final List<Set<String>> subSets = new ArrayList<Set<String>>(partitions);
        for (int i = 0; i < partitions; i++) {
            subSets.add(new HashSet<String>());
        }

        final int partitionSize = (setToPartition.size() % partitions == 0) ? setToPartition.size() / partitions : setToPartition.size() / partitions + 1;
        final Iterator<String> it = setToPartition.iterator();
        for (final Set<String> subSet : subSets) {
            while (subSet.size() < partitionSize && it.hasNext()) {
                subSet.add(it.next());
            }
        }

        return subSets;
    }

    private Set<String> getTokensToCount(String tokenFile) {
        final Set<String> uniqueTokens = new HashSet<String>();

        final Thread thread = new Thread(new ReadInputFileTask(tokenFile, uniqueTokens));
        thread.start();
        final long threadStartTime = System.currentTimeMillis();
        try {
            thread.join();
        } catch (InterruptedException e) {
            e.printStackTrace();
            System.exit(-1);
        }
        System.out.println("Preprocessing input file: " + (System.currentTimeMillis() - threadStartTime));

        return uniqueTokens;
    }

    private static List<String> getTextTokens(String textFile) {
        final List<String> textTokens = new LinkedList<String>();

        final Thread thread = new Thread(new ReadShakespeareTextTask(textFile, textTokens));
        thread.start();
        final long threadStartTime = System.currentTimeMillis();
        try {
            thread.join();
        } catch (InterruptedException e) {
            e.printStackTrace();
            System.exit(-1);
        }
        System.out.println("Preprocessing text file: " + (System.currentTimeMillis() - threadStartTime));

        return textTokens;
    }

    private void spawnAndRunWorkerThreads(List<String> tokensToScan, Set<String> uniqueTokens) {
        final long startTime = System.currentTimeMillis();
        
        final List<Set<String>> partitionedTokens = partitionSet(uniqueTokens, numberOfThreads);

        for (int i = 0; i < numberOfThreads; i++) {
            workerTasks[i] = new CountWordTask(tokensToScan, partitionedTokens.get(i));
            workerThreads[i] = new Thread(workerTasks[i]);
            workerThreads[i].start();
        }
        
        for (final Thread thread : workerThreads) {
            try {
                thread.join();
            } catch (InterruptedException e) {
                e.printStackTrace();
                System.exit(-1);
            }
        }
        System.out.println("Counting: " + (System.currentTimeMillis() - startTime));
    }

    private void writeFinalResult() {
        final long startTime = System.currentTimeMillis();
        
        final Map<String, Integer> result = new TreeMap<String, Integer>();
        for (final CountWordTask task : workerTasks) {
            result.putAll(task.getWordCount());
        }

        PrintWriter writer = null;
        try {
            final File outputFile = new File(outputFileName);
            if (outputFile.exists()) {
                outputFile.delete();
            }
            writer = new PrintWriter(outputFile);
            for (final Entry<String, Integer> wordCount : result.entrySet()) {
                writer.println(wordCount.getKey() + " #" + wordCount.getValue());
            }
        } catch (IOException e) {
            e.printStackTrace();
            System.exit(-1);
        } finally {
            writer.close();
        }
        
        System.out.println("Output: " + (System.currentTimeMillis() - startTime));
    }

    public static void main(String[] args) {
        if (args.length != 4) {
            System.out.println("Usage: java ShakespeareWordCountApp numOfThreads fullShakespeareTextFile inputTextFile outputFile");
            System.exit(0);
        }

        try {
            final long startTime = System.currentTimeMillis();

            final int numberOfThreads = Integer.valueOf(args[0]);
            final String fullShakespeareText = args[1];
            final String inputFile = args[2];
            final String ouputFile = args[3];

            new ShakespeareWordCountApp(numberOfThreads, fullShakespeareText, inputFile, ouputFile).run();
            System.out.println("Total: " + (System.currentTimeMillis() - startTime));
        } catch (Exception e) {
            e.printStackTrace();
            System.exit(-1);
        }
    }
}
